import numpy as np
import matplotlib.pyplot as plt

# Adam
class Adam:
    def __init__(self, params, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        self.params = params
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.m = [np.zeros_like(p) for p in params]
        self.v = [np.zeros_like(p) for p in params]
        self.t = 0
        
    def step(self, grads):
        self.t += 1
        for i in range(len(self.params)):
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grads[i]
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * (grads[i]**2)
            m_hat = self.m[i] / (1 - self.beta1**self.t)
            v_hat = self.v[i] / (1 - self.beta2**self.t)
            self.params[i] -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

# 数据
X = np.array([
    [1,2,3],
    [4,5,6],
    [7,8,9]
])
T = np.array([
    [10,10],
    [10,10]
])

W = np.random.randn(2,2)
b = 0.0

def conv2d(input, kernel, bias):
    h, w = input.shape
    kh, kw = kernel.shape
    out = np.zeros((h - kh + 1, w - kw + 1))
    for i in range(out.shape[0]):
        for j in range(out.shape[1]):
            out[i,j] = np.sum(input[i:i+kh,j:j+kw] * kernel) + bias
    return out

def mse(y, t):
    return np.mean((y - t)**2)

def mse_grad(y, t):
    N = y.size
    return 2*(y - t)/N

# adam 优化器
adam = Adam([W, b], lr=0.01)

loss_history = []

for epoch in range(1, 501):
    Y = conv2d(X, W, b)
    L = mse(Y, T)
    loss_history.append(L)
    
    if epoch % 50 == 0 or epoch <= 5:
        print(f"Epoch {epoch} Loss {L:.4f}")
        print("前向输出 Y:\n", Y)
    
    G = mse_grad(Y, T)
    
    grad_W = np.zeros_like(W)
    for m in range(W.shape[0]):
        for n in range(W.shape[1]):
            s = 0
            for i in range(G.shape[0]):
                for j in range(G.shape[1]):
                    s += G[i,j] * X[i+m,j+n]
            grad_W[m,n] = s
    
    grad_b = np.sum(G)
    
    adam.step([grad_W, grad_b])

plt.plot(loss_history)
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Loss with Bias + Adam")
plt.grid()
plt.show()
